from sympy import symbols, nroots, simplify, re, im
import re as pyRE
threshold = 1e-9

def solveEq(eq) -> dict:
  λ = symbols("λ")
  try:
    solutions = nroots(eq)
    solutions = [re(sol) if abs(im(sol)) < threshold else sol for sol in solutions]
    if all(s.is_real for s in solutions):
      allRealRoots = True  
      solutions = [float(str(s.evalf(10))) for s in solutions]
      solutions.sort(reverse=True)
    else:
      allRealRoots = False
      solutions = [str(s.evalf(3)) for s in solutions]
  except:
    print("solveWithNumpy:")
    solutions, allRealRoots = solveWithNumpy(eq)
  return {"roots":solutions,"allRealRoots":allRealRoots}
  
def solveWithNumpy(eq) -> dict:
  from numpy.polynomial import Polynomial as npPoly
  import numpy as np
  coef = []
  eq = "+" + eq if not eq.startswith("-") else eq
  l = pyRE.split(r"([\+\-])",eq)
  l = ["1*"+l[i] if l[i].startswith("λ") else l[i] for i in range(len(l))]
  l = list(filter(lambda x: x != "",l))
  l = [l[i]+l[i+1] for i in range(0, len(l)-1,2)]
  m = [len(pyRE.findall(r"\*λ",l)) for l in l]
  l= [float(l.replace(r"*λ","")) for l in l]
  for i in range(m[-1]+1):
    if i in m:
      coef.append(l[m.index(i)])
    else:
      coef.append(0)
  solutions = npPoly(coef,symbol="λ").roots()
  if (all(np.isreal(solutions))):
    allRealRoots = True
    solutions = [round(s.astype(float),10) for s in solutions]
    solutions.sort(reverse=True)
  else:
    allRealRoots = False
    solutions = np.array(solutions).astype(str)
  return solutions,allRealRoots

if __name__ == "__main__":
  solveEq(0)
